Fix loss weighting for short iterable evaluation batches - #50708

What does this PR do?

Use each observed evaluation batch size when expanding its mean loss before gathering.

With an unknown-length iterable dataset and accelerator_config={"dispatch_batches": False}, the final batch can be shorter than the configured batch size. Its mean loss is currently repeated for a full batch. Accelerate cannot trim this unknown-length loader by a known dataset remainder, so the tail is overweighted in both evaluate() and predict().

The existing configured-size fallback remains for inputs whose batch size cannot be inferred. Gathering, accumulation, and the reduction of model losses are unchanged.

Reproduction and validation

Native CPU validation: https://github.com/Afloat16/transformers/actions/runs/43406792756

• New regression in the existing iterable Trainer suite: original code has 2 failing and 10 passing subcases; the fix passes all 12 subcases. These cover actual evaluate()/predict(), an independent MSE reference, dispatch enabled/disabled, accumulation enabled/disabled, and batch sizes 1, 3, and 8 for seven examples.
• Four existing evaluation/prediction tests pass, including the existing iterable tests.
• A separate real workflow uses an unmodified default GPT2Config(), GPT2LMHeadModel, Trainer, and the same Hugging Face Dataset converted with to_iterable_dataset(). With batch size 3 and seven examples:

The GPT-2 check uses locally initialized weights and verifies evaluation aggregation; it is not a model-quality benchmark. Validation ran on CPU, without distributed or GPU claims.

Commands:

python -m pytest tests/trainer/test_trainer_data.py::TrainerDynamicShapesAndIterableTest::test_iterable_evaluation_loss_uses_actual_batch_size -q
python -m pytest tests/trainer/test_trainer_data.py::TrainerDynamicShapesAndIterableTest::test_evaluation_iterable_dataset tests/trainer/test_trainer_data.py::TrainerDynamicShapesAndIterableTest::test_predict_iterable_dataset tests/trainer/test_trainer_evaluation.py::TrainerEvaluationTest::test_evaluate tests/trainer/test_trainer_evaluation.py::TrainerEvaluationTest::test_predict -q
make style

Full make style passes, with 5,108 files unchanged.

Related work

This is a standalone reproduction. The existing batch-size fix in #16763 selected the configured total batch size for data parallelism. The Accelerate migration in #24028 removed the old final truncation of losses and delegated gathering to Accelerate; this patch instead makes each local loss expansion match the batch that produced it. Related evaluation-loop changes and current open/closed PRs were checked for the same fix.